from publicgood_experiment import *
from publicgood_agent import *
from publicgood_environment import *
from pybrain.tools.functions import sigmoid
import numpy
import matplotlib.pyplot as plt

class rule0:
    dim = 1
    def __call__(self, contributions, punishments):
        return numpy.sum(contributions)

class rule:
    dim = 4
    def __call__(self, contributions, punishments):
        """
            rule returns a n-player array with values from 0 to 1 with 0
            corresponding to players that do not follow the rule, and 1
            corresponding to players that do.

            allocations is everyone's allocations for that period, and
            punishments the correspond punishments for that period. The punishments
            should be a list of player's punishment lists.
        """
        rule.dim = 4
        assert contributions != None and punishments != None,  "no contributions or punishments"

        contribution_compliance = map(lambda cont: cont/20, contributions)
        ## Ci/y ... contribution divided by endowment. Always bounded 0 to 1.
        
        def punish_measure(ipunishments):
            return sum(map(lambda ijpunish, jcont: ijpunish / (21 - jcont), ipunishments, contributions))
        punishment_compliance = map(punish_measure, punishments)
        ## Sum of Pij/E-Cj ... the sum of the punishments by the differences in allocations.

        total_compliance = map(lambda x,y: 0.5*x + 0.5*sigmoid(y), contribution_compliance, punishment_compliance)
        ## Compliance is half of the contribution compliance, and half of the squished punishment compliance.

        print("Compliance measure")
        print(total_compliance)
        return total_compliance


class Run:
    experiments = {
     #  experiment  N       Observational vars                  punish          AI vars
     #              agents  allocations punishments other       effectivenes
        'exp0':     (4,     True,       False,      None,       0,              1, 1),  #This is a baseline experiment, no punishment. Just a simple public good game
        'exp0b':    (4,     False,      False,      rule0(),    0,              1, 1),  #More of the same, except agents can only see the sum of allocations
        'exp1':     (4,     True,       True,       None,       3,              1, 1),  #Agents can see everything.
        'exp2':     (4,     False,      False,      rule(),     3,              1, 1),  #Agents "understand the common institution"
        'bug1':     (4,     True,       True,       None,       3,              0, 0)   #To-do: get the NFQ learning agent working.
    }

    def __init__(self, verbosity = 1, graph = True, write_file=None, measures = []):
     
        self.measures = measures
        self.verbosity = verbosity
        if graph:
            self.graphing = True
    
    def __call__(self, run_count=150, exp='exp0', repeat=100):
        self.general_setup(run_count, repeat, *Run.experiments[exp])

    def general_setup(self, run_count, repeat, nplayers, return_allocations, return_punishments, return_rule, punish_eff, alloc_type, punish_type):
        env = PublicGoodEnvironment(nplayers, epsilon=punish_eff)
        #The environment itself is composed of sub-environments to work better with pyBrain classes.
        env.set_sub_environments([return_allocations, return_punishments, return_rule],[return_allocations, return_punishments, return_rule])
        #The experiment creates its own tasks and agents, which are nested nplayer length arrays with two sub-elements for both the punish and allocate tasks.

        #Verbosity=True lets us see the player's observations.
        def allos(idnumber, *args):
            return args[0][idnumber]
        def puns(idnumber, *args):
            return args[1][idnumber].mean()
        self.measures.insert(0, allos)
        if env.epsilon != 0:
            self.measures.insert(1, puns)
        data = numpy.zeros((repeat, len(self.measures), run_count))
        for i in range(repeat):
            env.reset()
            exp = PublicGoodExperiment(env, alloc_type, punish_type)
            data[i,...,...] = exp.doEpisodes(run_count, self.verbosity, self.measures).mean(axis=1)

        if self.verbosity>=1:
            print("Data:")
            print(data)
        data = data.mean(axis=0)
        if self.verbosity>=1:
            print("Avg Data")
            print(data)
        if self.graphing:
            pts = ['bs', 'rs', 'yt', 'go']
            for i, measure in enumerate(data):
                if self.verbosity==2:
                    print("Measure data:")
                    print(measure)
                plt.plot(measure, pts[i])
            plt.show()




